from LeetCoder import *


#1000000
class Solution:

    def punishmentNumber(self, n: int) -> int:

        def g(nums, total, index, num):
            if index == len(nums):
                return num == total
            c = 0
            for i in range(index, len(nums)):
                c = c * 10 + nums[i]
                if g(nums, c + total, i + 1, num):
                    return True
            return False

        def f(num):
            s = list(map(int, str(num * num)))
            if g(s, 0, 0, num):
                return True

        ans = 0
        for i in range(1, n + 1):
            if f(i):
                ans += i * i
        return ans


s = Solution()
print(s.punishmentNumber(10))
print(s.punishmentNumber(37))
